"""Test suite to check index implementations.

Standard tests for the DocumentIndex abstraction

We don't recommend implementing externally managed DocumentIndex abstractions at this
time.
"""

import inspect
import uuid
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator, Generator

import pytest
from langchain_core.documents import Document
from langchain_core.indexing.base import DocumentIndex


class DocumentIndexerTestSuite(ABC):
    """Test suite for checking the read-write of a document index.

    Implementers should subclass this test suite and provide a fixture
    that returns an empty index for each test.
    """

    @abstractmethod
    @pytest.fixture
    def index(self) -> Generator[DocumentIndex, None, None]:
        """Get the index."""

    def test_upsert_documents_has_no_ids(self, index: DocumentIndex) -> None:
        """Verify that there is no parameter called IDs in upsert."""
        signature = inspect.signature(index.upsert)
        assert "ids" not in signature.parameters

    def test_upsert_no_ids(self, index: DocumentIndex) -> None:
        """Upsert works with documents that do not have IDs.

        At the moment, the ID field in documents is optional.
        """
        documents = [
            Document(page_content="foo", metadata={"id": 1}),
            Document(page_content="bar", metadata={"id": 2}),
        ]
        response = index.upsert(documents)
        ids = sorted(response["succeeded"])

        # Ordering is not guaranteed, need to test carefully
        documents = index.get(ids)
        sorted_documents = sorted(documents, key=lambda x: x.id or "")

        if sorted_documents[0].page_content == "bar":
            assert sorted_documents[0] == Document(
                page_content="bar", metadata={"id": 2}, id=ids[0]
            )
            assert sorted_documents[1] == Document(
                page_content="foo", metadata={"id": 1}, id=ids[1]
            )
        else:
            assert sorted_documents[0] == Document(
                page_content="foo", metadata={"id": 1}, id=ids[0]
            )
            assert sorted_documents[1] == Document(
                page_content="bar", metadata={"id": 2}, id=ids[1]
            )

    def test_upsert_some_ids(self, index: DocumentIndex) -> None:
        """Test an upsert where some docs have IDs and some don't."""
        foo_uuid = str(uuid.UUID(int=7))
        documents = [
            Document(id=foo_uuid, page_content="foo", metadata={"id": 1}),
            Document(page_content="bar", metadata={"id": 2}),
        ]
        response = index.upsert(documents)
        ids = response["succeeded"]
        other_id = next(iter(set(ids) - {foo_uuid}))
        assert response["failed"] == []
        assert foo_uuid in ids
        # Ordering is not guaranteed, so we use a set.
        documents = index.get(ids)
        first_doc = documents[0]
        if first_doc.id == foo_uuid:
            assert documents == [
                Document(page_content="foo", metadata={"id": 1}, id=foo_uuid),
                Document(page_content="bar", metadata={"id": 2}, id=other_id),
            ]
        else:
            assert documents == [
                Document(page_content="bar", metadata={"id": 2}, id=other_id),
                Document(page_content="foo", metadata={"id": 1}, id=foo_uuid),
            ]

    def test_upsert_overwrites(self, index: DocumentIndex) -> None:
        """Test that upsert overwrites existing content."""
        foo_uuid = str(uuid.UUID(int=7))
        documents = [
            Document(id=foo_uuid, page_content="foo", metadata={"bar": 1}),
        ]
        response = index.upsert(documents)
        ids = response["succeeded"]
        assert response["failed"] == []

        assert index.get(ids) == [
            Document(page_content="foo", metadata={"bar": 1}, id=foo_uuid),
        ]

        # Now let's overwrite foo
        index.upsert([Document(id=foo_uuid, page_content="foo2", metadata={"meow": 2})])
        documents = index.get([foo_uuid])
        assert documents == [
            Document(page_content="foo2", metadata={"meow": 2}, id=foo_uuid)
        ]

    def test_delete_missing_docs(self, index: DocumentIndex) -> None:
        """Verify that we can delete docs that aren't there."""
        assert index.get(["1"]) == []  # Should be empty.

        delete_response = index.delete(["1"])
        if "num_deleted" in delete_response:
            assert delete_response["num_deleted"] == 0

        if "num_failed" in delete_response:
            # Deleting a missing an ID is **not** failure!!
            assert delete_response["num_failed"] == 0

        if "succeeded" in delete_response:
            # There was nothing to delete!
            assert delete_response["succeeded"] == []

        if "failed" in delete_response:
            # Nothing should have failed
            assert delete_response["failed"] == []

    def test_delete_semantics(self, index: DocumentIndex) -> None:
        """Test deletion of content has appropriate semantics."""
        # Let's index a document first.
        foo_uuid = str(uuid.UUID(int=7))
        upsert_response = index.upsert(
            [Document(id=foo_uuid, page_content="foo", metadata={})]
        )
        assert upsert_response == {"succeeded": [foo_uuid], "failed": []}

        delete_response = index.delete(["missing_id", foo_uuid])

        if "num_deleted" in delete_response:
            assert delete_response["num_deleted"] == 1

        if "num_failed" in delete_response:
            # Deleting a missing an ID is **not** failure!!
            assert delete_response["num_failed"] == 0

        if "succeeded" in delete_response:
            # There was nothing to delete!
            assert delete_response["succeeded"] == [foo_uuid]

        if "failed" in delete_response:
            # Nothing should have failed
            assert delete_response["failed"] == []

    def test_bulk_delete(self, index: DocumentIndex) -> None:
        """Test that we can delete several documents at once."""
        documents = [
            Document(id="1", page_content="foo", metadata={"id": 1}),
            Document(id="2", page_content="bar", metadata={"id": 2}),
            Document(id="3", page_content="baz", metadata={"id": 3}),
        ]

        index.upsert(documents)
        index.delete(["1", "2"])
        assert index.get(["1", "2", "3"]) == [
            Document(page_content="baz", metadata={"id": 3}, id="3")
        ]

    def test_delete_no_args(self, index: DocumentIndex) -> None:
        """Test delete with no args raises ValueError."""
        with pytest.raises(ValueError):  # noqa: PT011
            index.delete()

    def test_delete_missing_content(self, index: DocumentIndex) -> None:
        """Deleting missing content should not raise an exception."""
        index.delete(["1"])
        index.delete(["1", "2", "3"])

    def test_get_with_missing_ids(self, index: DocumentIndex) -> None:
        """Test get with missing IDs."""
        documents = [
            Document(id="1", page_content="foo", metadata={"id": 1}),
            Document(id="2", page_content="bar", metadata={"id": 2}),
        ]
        upsert_response = index.upsert(documents)
        assert upsert_response == {
            "succeeded": ["1", "2"],
            "failed": [],
        }
        retrieved_documents = index.get(["1", "2", "3", "4"])
        # The ordering is not guaranteed, so we use a set.
        assert sorted(retrieved_documents, key=lambda x: x.id or "") == [
            Document(page_content="foo", metadata={"id": 1}, id="1"),
            Document(page_content="bar", metadata={"id": 2}, id="2"),
        ]

    def test_get_missing(self, index: DocumentIndex) -> None:
        """Test get by IDs with missing IDs."""
        # This should not raise an exception
        documents = index.get(["1", "2", "3"])
        assert documents == []


class AsyncDocumentIndexTestSuite(ABC):
    """Test suite for checking the read-write of a document index.

    Implementers should subclass this test suite and provide a fixture
    that returns an empty index for each test.
    """

    @abstractmethod
    @pytest.fixture
    async def index(self) -> AsyncGenerator[DocumentIndex, None]:
        """Get the index."""

    async def test_upsert_documents_has_no_ids(self, index: DocumentIndex) -> None:
        """Verify that there is not parameter called IDs in upsert."""
        signature = inspect.signature(index.upsert)
        assert "ids" not in signature.parameters

    async def test_upsert_no_ids(self, index: DocumentIndex) -> None:
        """Upsert works with documents that do not have IDs.

        At the moment, the ID field in documents is optional.
        """
        documents = [
            Document(page_content="foo", metadata={"id": 1}),
            Document(page_content="bar", metadata={"id": 2}),
        ]
        response = await index.aupsert(documents)
        ids = sorted(response["succeeded"])

        # Ordering is not guaranteed, need to test carefully
        documents = await index.aget(ids)
        sorted_documents = sorted(documents, key=lambda x: x.id or "")

        if sorted_documents[0].page_content == "bar":
            assert sorted_documents[0] == Document(
                page_content="bar", metadata={"id": 2}, id=ids[0]
            )
            assert sorted_documents[1] == Document(
                page_content="foo", metadata={"id": 1}, id=ids[1]
            )
        else:
            assert sorted_documents[0] == Document(
                page_content="foo", metadata={"id": 1}, id=ids[0]
            )
            assert sorted_documents[1] == Document(
                page_content="bar", metadata={"id": 2}, id=ids[1]
            )

    async def test_upsert_some_ids(self, index: DocumentIndex) -> None:
        """Test an upsert where some docs have IDs and some don't."""
        foo_uuid = str(uuid.UUID(int=7))
        documents = [
            Document(id=foo_uuid, page_content="foo", metadata={"id": 1}),
            Document(page_content="bar", metadata={"id": 2}),
        ]
        response = await index.aupsert(documents)
        ids = response["succeeded"]
        other_id = next(iter(set(ids) - {foo_uuid}))
        assert response["failed"] == []
        assert foo_uuid in ids
        # Ordering is not guaranteed, so we use a set.
        documents = await index.aget(ids)
        first_doc = documents[0]
        if first_doc.id == foo_uuid:
            assert documents == [
                Document(page_content="foo", metadata={"id": 1}, id=foo_uuid),
                Document(page_content="bar", metadata={"id": 2}, id=other_id),
            ]
        else:
            assert documents == [
                Document(page_content="bar", metadata={"id": 2}, id=other_id),
                Document(page_content="foo", metadata={"id": 1}, id=foo_uuid),
            ]

    async def test_upsert_overwrites(self, index: DocumentIndex) -> None:
        """Test that upsert overwrites existing content."""
        foo_uuid = str(uuid.UUID(int=7))
        documents = [
            Document(id=foo_uuid, page_content="foo", metadata={"bar": 1}),
        ]
        response = await index.aupsert(documents)
        ids = response["succeeded"]
        assert response["failed"] == []

        assert await index.aget(ids) == [
            Document(page_content="foo", metadata={"bar": 1}, id=foo_uuid),
        ]

        # Now let's overwrite foo
        await index.aupsert(
            [Document(id=foo_uuid, page_content="foo2", metadata={"meow": 2})]
        )
        documents = await index.aget([foo_uuid])
        assert documents == [
            Document(page_content="foo2", metadata={"meow": 2}, id=foo_uuid)
        ]

    async def test_delete_missing_docs(self, index: DocumentIndex) -> None:
        """Verify that we can delete docs that aren't there."""
        assert await index.aget(["1"]) == []  # Should be empty.

        delete_response = await index.adelete(["1"])
        if "num_deleted" in delete_response:
            assert delete_response["num_deleted"] == 0

        if "num_failed" in delete_response:
            # Deleting a missing an ID is **not** failure!!
            assert delete_response["num_failed"] == 0

        if "succeeded" in delete_response:
            # There was nothing to delete!
            assert delete_response["succeeded"] == []

        if "failed" in delete_response:
            # Nothing should have failed
            assert delete_response["failed"] == []

    async def test_delete_semantics(self, index: DocumentIndex) -> None:
        """Test deletion of content has appropriate semantics."""
        # Let's index a document first.
        foo_uuid = str(uuid.UUID(int=7))
        upsert_response = await index.aupsert(
            [Document(id=foo_uuid, page_content="foo", metadata={})]
        )
        assert upsert_response == {"succeeded": [foo_uuid], "failed": []}

        delete_response = await index.adelete(["missing_id", foo_uuid])

        if "num_deleted" in delete_response:
            assert delete_response["num_deleted"] == 1

        if "num_failed" in delete_response:
            # Deleting a missing an ID is **not** failure!!
            assert delete_response["num_failed"] == 0

        if "succeeded" in delete_response:
            # There was nothing to delete!
            assert delete_response["succeeded"] == [foo_uuid]

        if "failed" in delete_response:
            # Nothing should have failed
            assert delete_response["failed"] == []

    async def test_bulk_delete(self, index: DocumentIndex) -> None:
        """Test that we can delete several documents at once."""
        documents = [
            Document(id="1", page_content="foo", metadata={"id": 1}),
            Document(id="2", page_content="bar", metadata={"id": 2}),
            Document(id="3", page_content="baz", metadata={"id": 3}),
        ]

        await index.aupsert(documents)
        await index.adelete(["1", "2"])
        assert await index.aget(["1", "2", "3"]) == [
            Document(page_content="baz", metadata={"id": 3}, id="3")
        ]

    async def test_delete_no_args(self, index: DocumentIndex) -> None:
        """Test delete with no args raises ValueError."""
        with pytest.raises(ValueError):  # noqa: PT011
            await index.adelete()

    async def test_delete_missing_content(self, index: DocumentIndex) -> None:
        """Deleting missing content should not raise an exception."""
        await index.adelete(["1"])
        await index.adelete(["1", "2", "3"])

    async def test_get_with_missing_ids(self, index: DocumentIndex) -> None:
        """Test get with missing IDs."""
        documents = [
            Document(id="1", page_content="foo", metadata={"id": 1}),
            Document(id="2", page_content="bar", metadata={"id": 2}),
        ]
        upsert_response = await index.aupsert(documents)
        assert upsert_response == {
            "succeeded": ["1", "2"],
            "failed": [],
        }
        retrieved_documents = await index.aget(["1", "2", "3", "4"])
        # The ordering is not guaranteed, so we use a set.
        assert sorted(retrieved_documents, key=lambda x: x.id or "") == [
            Document(page_content="foo", metadata={"id": 1}, id="1"),
            Document(page_content="bar", metadata={"id": 2}, id="2"),
        ]

    async def test_get_missing(self, index: DocumentIndex) -> None:
        """Test get by IDs with missing IDs."""
        # This should not raise an exception
        documents = await index.aget(["1", "2", "3"])
        assert documents == []
